import os
import sqlite3
import hashlib
from pathlib import Path
from datetime import datetime
import sys
import argparse
from PIL import Image
from PIL.ExifTags import TAGS
from tqdm import tqdm

def is_photo_file(filename):
    """判断给定的文件是否为视频文件"""
    image_extensions = {'.jpg', '.jpeg', '.png','.aae','.heic','.thm', '.gif', '.bmp', '.tiff', '.webp'}
    return Path(filename).suffix.lower() in image_extensions

# 注册datetime适配器和转换器
def adapt_datetime(ts):
    return ts.strftime('%Y-%m-%d %H:%M:%S')

def convert_datetime(s):
    return datetime.strptime(s.decode('utf-8'), '%Y-%m-%d %H:%M:%S')

sqlite3.register_adapter(datetime, adapt_datetime)
sqlite3.register_converter("datetime", convert_datetime)

def check_and_create_db(db_path):
    conn = sqlite3.connect(db_path, detect_types=sqlite3.PARSE_DECLTYPES)
    cursor = conn.cursor()
    
    # 检查photo_hash表是否存在
    cursor.execute("SELECT name FROM sqlite_master WHERE type='table' AND name='photo_hash';")
    if cursor.fetchone() is None:
        print("请先运行 photo_scan.py 创建必要的数据库表结构(photo_hash)。")
        sys.exit(1)
    
    # 创建photo_move表结构
    cursor.execute('''
        CREATE TABLE IF NOT EXISTS photo_move (
            id INTEGER PRIMARY KEY,
            file_path TEXT NOT NULL,
            file_name TEXT NOT NULL,
            file_time datetime NOT NULL,
            file_hash TEXT NOT NULL,
            aim_zip TEXT NOT NULL,
            deal_flag INTEGER DEFAULT 0,
            create_time datetime NOT NULL
        )
    ''')
    conn.commit()
    return conn

def calculate_file_hash(file_path):
    """计算给定文件内容的SHA256哈希值"""
    sha256_hash = hashlib.sha256()
    with open(file_path, "rb") as f:
        for byte_block in iter(lambda: f.read(4096), b""):
            sha256_hash.update(byte_block)
    return sha256_hash.hexdigest()

def get_photo_time(file_path):
    """获取文件的拍摄时间或修改时间"""
    try:
        image = Image.open(file_path)
        exif_data = image._getexif()
        if exif_data is not None:
            for tag, value in exif_data.items():
                tag_name = TAGS.get(tag, tag)
                if tag_name == 'DateTimeOriginal':
                    return datetime.strptime(value, '%Y:%m:%d %H:%M:%S')
    except Exception as e:
        pass
    return datetime.fromtimestamp(os.path.getmtime(file_path))

def determine_aim_zip(file_time):
    """根据文件时间确定目标ZIP文件名"""
    year = file_time.year
    quarter = (file_time.month - 1) // 3 + 1
    return f"相片{year}/iPhone{year}Q{quarter}.zip"

def scan_and_store_files(directory, db_conn):
    cursor = db_conn.cursor()
    
    base_dir = Path(directory)
    
    # 计算总文件数，用于进度条
    total_files = sum(len(files) for _, _, files in os.walk(base_dir))
    
    with tqdm(total=total_files, desc='Processing files', unit='file') as pbar:
        for root, dirs, files in os.walk(base_dir):
            for file_name in files:
                pbar.update(1)
                file_path = Path(root) / file_name
                if not is_photo_file(file_name) :
                    continue
                
                relative_path = str(file_path.relative_to(base_dir).parent)
                
                # 计算文件哈希值
                file_hash = calculate_file_hash(file_path)
                
                # 检查哈希是否已存在于photo_hash表中
                cursor.execute("SELECT COUNT(*) FROM photo_hash WHERE file_hash=?", (file_hash,))
                if cursor.fetchone()[0] > 0:
                    continue
                
                cursor.execute("SELECT COUNT(*) FROM photo_move WHERE file_hash=?", (file_hash,))
                if cursor.fetchone()[0] > 0:
                    continue

                # 获取文件时间
                file_time = get_photo_time(file_path)
                
                # 确定目标ZIP文件名
                aim_zip = determine_aim_zip(file_time)
                
                # 插入记录到photo_move表
                current_time = datetime.now()
                cursor.execute('''
                    INSERT INTO photo_move (file_path, file_name, file_time, file_hash, aim_zip, create_time)
                    VALUES (?, ?, ?, ?, ?, ?)
                ''', (
                    relative_path,
                    file_name,
                    file_time,
                    file_hash,
                    aim_zip,
                    current_time
                ))
                print(f"已添加文件: {file_path}")
        

      
    
    db_conn.commit()

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Scan directory and store file information into the database.")
    parser.add_argument('--db', required=True, help="Directory containing the SQLite database file (file_hashes.db)")
    parser.add_argument('--scan-dir', required=True, help="Directory to scan for files")

    args = parser.parse_args()

    db_directory = args.db
    scan_directory = args.scan_dir

    if not os.path.isdir(db_directory):
        print(f"错误: {db_directory} 不是一个有效的目录")
        sys.exit(1)

    if not os.path.isdir(scan_directory):
        print(f"错误: {scan_directory} 不是一个有效的目录")
        sys.exit(1)

    db_path = str(Path(db_directory) / 'file_hashes.db')
    conn = check_and_create_db(db_path)
    scan_and_store_files(scan_directory, conn)
    conn.close()